{-# LANGUAGE ForeignFunctionInterface #-}
module Hexena.Cmds.GenerateReport (run) where

import Data.Array.Unboxed
import Data.Binary.Get
import qualified Data.ByteString as B
import qualified Data.ByteString.Lazy as BL
import Data.IFF
import Data.List (find, isPrefixOf, isSuffixOf)
import qualified Data.IntMap as IM
import qualified Data.Map as M
import Data.Maybe
import Data.Version (showVersion)
import Data.Word
import Foreign (Ptr, withForeignPtr)
import Foreign.C.String
import Graphics.UI.SDL
import System.Cmd (system)
import System.Directory (createDirectoryIfMissing)
import System.Exit
import Text.Printf

--import Configuration -- Using Paths_hexena as generated by cabal
import Paths_hexena (getDataFileName, version)

import Hexena.Chunks
import Hexena.Text.DrawGraph
import Hexena.Text.Hexdump
import Hexena.MZHeader
import Hexena.PEHeader
import Hexena.Reporting
import Hexena.Math.Statistics (Histogram, calculateE, scaleHist)
import qualified Hexena.Modules.BlockEntropy
import Hexena.Internal.SVMWitchcraft

-- TODO use Paths_Hexena getDataFileName (FilePath -> IO FilePath)

histogramGraphName :: String
histogramGraphName = "1-gram_histogram"

logoFilename :: String
logoFilename = "hexena.eps"

printMZ :: BL.ByteString -> String
printMZ x = "\\begin{verbatim}\n" ++ unlines (compactHexify x) ++ "\\end{verbatim}"

--extern "C" void writeXPM(SDL_Surface *surf, const char *fname) {
foreign import ccall unsafe "writeXPM" writeXPM :: Ptr SurfaceStruct -> CString -> IO ()

maxHeight :: Int
maxHeight = 64

_MAL_TEMP :: String
_MAL_TEMP = "malware_report.template.tex"

displayHGraph :: String -> UArray Word8 Int -> UArray Word8 Int -> IO ()
displayHGraph f h1 h2 = withInit [InitVideo] $ do
  display <- setVideoMode 270 (2 * fromIntegral maxHeight) 32 []
  let disform = surfaceGetPixelFormat display
  setCaption f f
  _ <- fillRect display Nothing =<< mapRGB disform 0xff 0xff 0xff
  _ <- fillRect display (Just $ Rect (7+0x20) 0 0x10 270) =<< mapRGB disform 0xDA 0xD3 0xD3
  _ <- fillRect display (Just $ Rect (7+0x30) 0 0x10 270) =<< mapRGB disform 0xCA 0xC3 0xC3
  _ <- fillRect display (Just $ Rect (7+0x40) 0 0x40 270) =<< mapRGB disform 0xBA 0xB6 0xB6
  drawHist display 0x5555a3ff 0x000080ff 7 64 $ elems h1'
  drawHist display 0x55A555ff 0x006010ff 7 66 $ elems h2'
  Graphics.UI.SDL.flip display
  withForeignPtr display $ withCString (f ++ ".xpm") . writeXPM
  -- 1-gram_histogram.eps3
      where h1' = scaleHist h1 id
            h2' = scaleHist h2 negate

getHashes :: [T] -> (String, String)
getHashes iff = let formHash = parts $ chunk $ fromJust $ find (isFormId (chunkIdFromString "HASH")) iff
                    readChunk x = B.unpack $ content $ chunk $ fromJust $ find ((==) x . chunkIdToString . chunkId) formHash
                    md5 = readChunk "MD5 "
                    sha1 = readChunk "SHA1"
                    hexPrint = foldr ((++) . printf "%02x") []
                in (hexPrint md5, hexPrint sha1)

-- | Generate the report entry for the lfanew part in the report.
-- mz/lfanew
reportLfanew :: IFFMap -> String
reportLfanew iff = txt ++ if lf `elem` [0xd8,0xe0..0xf8] then ok else strange
    where mz = content $ fromJust $ M.lookup "MZHD" iff
          lf = runGet getWord32be $ BL.fromChunks [B.drop 60 mz]
          txt = printf "It is located at position \\$%08x. " lf
          ok = "The PE header can be found at a normal position."
          strange = "The PE header is located at a non-standard position in the file."

tranRep :: IFFMap -> [T] -> [(String, String)] -> [String] -> [String]
tranRep iff rawiff other = map f
    where hdr = "%-- hexena: "
          peheader = loadPEHD iff
          (md5, sha1) = getHashes rawiff
          numbOfSect = peNumberOfSections peheader
          otherMap = M.fromList other
          f x | hdr `isPrefixOf` x = f' $ drop (length hdr) x
              | otherwise = x
          f' "pe/number_of_sections" = if numbOfSect `elem` [4,5] then "The number of sections found in this file is typical for a normal executable." else "The file contains an unusual number of sections."
          f' "hash/md5" = "MD5 & " ++ md5 ++ "\\\\"
          f' "hash/sha1" = "SHA1 & " ++ sha1 ++ "\\\\"
          f' "mz/lfanew" = reportLfanew iff
          f' "version" = showVersion version ++ "."
          f' x = f'' (M.lookup x otherMap)
              where f'' Nothing = hdr ++ "--#" ++ x
                    f'' (Just x') = x'


transformReport :: IFFMap -> [T] -> BL.ByteString -> Histogram -> Histogram -> [String] -> String
transformReport iff _ mz h1 h1C ts = unlines $ map fTrans ts
    where entr1 = calculateE h1
          entr1C = calculateE h1C
          mz' = mzDistance $ BL.unpack mz
          newcom x = "\\newcommand{\\hexena" ++ x ++ "}{"
          fTrans :: String -> String
          fTrans l | lFilesize `isPrefixOf` l = lFilesize ++ xFilesize ++ "}"
                   | lEntropy `isPrefixOf` l = lEntropy ++ show entr1 ++ "}"
                   | lREntropy `isPrefixOf` l = lREntropy ++ show entr1C ++ "}"
                   | lEntropyQuotient `isPrefixOf` l = lEntropyQuotient ++ show (entr1/entr1C) ++ "}"
                   | lMZHeaderNorm `isPrefixOf` l = lMZHeaderNorm ++ show mz' ++ "}"
                   | otherwise = l
          lFilesize = newcom "Filesize"
          lEntropy = newcom "Entropy"
          lREntropy = newcom "REntropy"
          lEntropyQuotient = newcom "EntropyQuotient"
          lMZHeaderNorm = newcom "MZHeaderNorm"
          xFilesize = show $ runGet getWord32be $ BL.fromChunks [content $ fromJust $ M.lookup "MWHD" iff]

classififySVM :: Maybe Model -> IFFMap -> IO (String, String)
classififySVM Nothing _ = error "model file cannot be found"
classififySVM (Just model) iff = let (_, lfanew) = loadMZHD $ fromJust $ M.lookup "MZHD" iff
                                     mz' = content $ fromJust $ M.lookup "MZRW" iff
                                     mzDist = mzDistanceN $ B.unpack mz'
                                     (hist, histC) = loadFORMHIST iff
                                     entr1 = calculateE hist
                                     entr1C = calculateE histC
                                     numberOfSections pe = fromIntegral (min sects 32) / 32.0
                                         where sects = peNumberOfSections pe
                                     lfanewReduced = if lfanew < 0x800 then fromIntegral lfanew / 0x800 else 1.0
                                     vals = [ numberOfSections $ loadPEHD iff
                                            , mzDist
                                            , lfanewReduced
                                            , entr1 / 8.0
                                            , entr1C / 8.0
                                            ]
                                     vect = IM.fromList $ zip [1..] vals
                                 in do r <- predict model vect
                                       print (r,vect)
                                       return ("classification/svm/vostro7", if r < 0 then "malicious" else "benign")

getAnalysedModules :: [T] -> [([T] -> IO (Maybe Report))] -> IO [Report]
getAnalysedModules iff fun = do rep <- mapM (\x -> x iff) fun
                                return $ map fromJust $ filter isJust rep
                                

writeModuleTex :: String -> Report -> IO ()
writeModuleTex dir r | null $ templateContent r = return ()
                     | otherwise = do writeFile (dir ++ reportName r ++ ".tex") $ templateContent r
                                      mapM_ (\(c, p) -> writeFile (dir ++ p) c) $ additionalFiles r

runner :: String -> String -> IFFMap -> [T] -> IO ()
runner template dir iff rawiff = let
    (hist, histC) = loadFORMHIST iff
    mzraw = BL.fromChunks [content $ fromJust $ M.lookup "MZRW" iff]
    peheader = loadPEHD iff
    templateLines = lines template
    modulesToRun = [(Hexena.Modules.BlockEntropy.reportFromIFF)]
  in do print iff
        model <- getDataFileName "svm.vostro7.model" >>= loadModel
        svmClass <- classififySVM model iff
        reports <- getAnalysedModules rawiff modulesToRun
        let lines' = tranRep iff rawiff ([svmClass] ++ concatMap globalReplacements reports) templateLines
        mapM_ (writeModuleTex dir) reports
        writeFile (dir ++ "malware_report.tex") $ transformReport iff rawiff mzraw hist histC lines'
        writeFile (dir ++ "mzheader.dump.tex") $ printMZ mzraw
        displayHGraph (dir ++ histogramGraphName) hist histC
        writeFile (dir ++ "peheader.table.tex") $ unlines $ map (\(k,v) -> printf "%s & \\hexval{%X} & %d \\\\" k v v) $ peheader2KeyValList peheader
        writeFile (dir ++ "peoptionalheader.table.tex") $ unlines $ map (\(k,v) -> printf "%s & \\hexval{%X} & %d \\\\" k v v) $ peoheader2KeyValList $ peOptionalHeader peheader
        writeFile (dir ++ "annotation.tex") $ loadANNO iff
        putStrLn "OK!"

{-main :: IO ()
main = do
  fname <- head `fmap` getArgs
  template <- tryLoad "malware_report.template.tex"
  (dat, rawiff) <- loadMalwIFF' fname
  putStrLn $ "Generating report for file: " ++ fname
  main2 template fname dat rawiff-}

appendReport :: String -> String
appendReport x | ".hexena" `isSuffixOf` x = take (length x - 7) x ++ suff
               | otherwise = x ++ suff
    where suff = ".report/"

convertHistogram :: String -> IO ()
convertHistogram n = let xn = n ++ ".xpm"
                         en = n ++ ".eps3"
                     in do x <- system $ printf "convert %s %s" xn en
                           case x of
                             ExitSuccess -> return ()
                             ExitFailure y -> error $ printf "external progam convert returned with error code %d" y

run :: FilePath -> IO ()
run fn = do
  let dir = appendReport fn
  createDirectoryIfMissing False dir
  template <- getDataFileName _MAL_TEMP >>= readFile
  (dat, rawiff) <- loadMalwIFF' fn
  putStrLn $ "Generating report for file: " ++ fn
  runner template dir dat rawiff
  convertHistogram $ dir ++ histogramGraphName
  getDataFileName logoFilename >>= readFile >>= writeFile (dir ++ logoFilename)
